import logging
from functools import partial

import numpy as np

from datasets.template_match_dataset import TemplateMatchDataset, MaskTemplateMatchDataset
from tools.QXSLAB_tools import QXSLAB_tools
from tools.sentinel12_tools import sentinel12_tools
from tools.os_tools import os_tools
from tools.dataset_tools import pm_mask_tools


class Sentinel12PairDataset(TemplateMatchDataset):


    def __init__(self, config=None):
        super(Sentinel12PairDataset, self).__init__(
            config, tools=sentinel12_tools, name="Sentinel12"
        )


class OSPairDataset(TemplateMatchDataset):


    def __init__(self, config=None):
        super(OSPairDataset, self).__init__(
            config, tools=os_tools, name="OS"
        )


class QXSLABPairDataset(TemplateMatchDataset):


    def __init__(self, config=None):
        super(QXSLABPairDataset, self).__init__(
            config, tools=QXSLAB_tools, name="QXSLAB"
        )


class OSMaskTMDataset(MaskTemplateMatchDataset):


    def __init__(self, config=None):
        super(OSMaskTMDataset, self).__init__(
            config, tools=pm_mask_tools, name="OS"
        )


# test datasets
if __name__ == "__main__":
    import cv2 as cv

    # dset_kwt = "os"
    # list_file = "E:/workspace/SOMatch/json/os_dataset/256_train.json"
    # dset_obj = OSPairDataset

    # dset_kwt = "sen12"
    # list_file = "E:/workspace/SOMatch/json/sen12_summer_harris/HARRIS_o1s1.json"
    # dset_obj = Sentinel12PairDataset

    dset_kwt = "QXSLAB"
    list_file = "E:/workspace/SOMatch/json/QXSLAB_SAROPT/QingDao.json"
    dset_obj = QXSLABPairDataset

    class crop:
        crop_size = 208
        shift_x = 15
        shift_y = 15
        search_rad = 32
        search_domain = 'A'

    class config:

        def __init__(self):
            self.single_domain = ""
            self.cache_dir = f"cache/{dset_kwt}/train"
            self.cache_size = 5000
            self.base_dir = None
            self.list_file = list_file
            self.augment = {
                "fliplr": 0,
                "flipud": 0,
                "scale": 0,
                "scale_px": (1.0, 1.0),
                "translate": 0,
                "translate_perc": (0.0, 0.0),
                "rotate": 0,
                "rotate_angle": (-5, 5)
            }
            self.template = crop()

    cfg = config()

    ppd = dset_obj(cfg)

    (img0, img1), (gt_map, _) = ppd[10]
    img0 = img0.numpy()
    img1 = img1.numpy()
    gt_map = gt_map.numpy()
    img0 = np.transpose(img0, [1, 2, 0])
    img1 = np.transpose(img1, [1, 2, 0])

    cv.imwrite(f"image/dataset_test/{dset_kwt}_sar.png", img0*255)
    cv.imwrite(f"image/dataset_test/{dset_kwt}_opt.png", img1*255)
    cv.imwrite(f"image/dataset_test/{dset_kwt}_map.png", gt_map*255)

